"""
数据处理工具模块

用于处理和分析子通道计算结果
"""

import numpy as np
from .. import config as cfg

class DataProcessor:
    """数据处理类"""
    
    @staticmethod
    def calculate_average_values(channels):
        """
        计算各个物理量的平均值
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        dict, 包含各个物理量平均值的字典
        """
        n_channels = len(channels)
        
        # 初始化结果字典
        averages = {
            'alpha_l': np.zeros(cfg.AXIAL_NODES),
            'alpha_v': np.zeros(cfg.AXIAL_NODES),
            'u_l': np.zeros(cfg.AXIAL_NODES),
            'u_v': np.zeros(cfg.AXIAL_NODES),
            'T_l': np.zeros(cfg.AXIAL_NODES),
            'T_v': np.zeros(cfg.AXIAL_NODES),
            'pressure': np.zeros(cfg.AXIAL_NODES),
            'T_w': np.zeros(cfg.AXIAL_NODES)
        }
        
        # 计算各物理量的平均值
        for channel in channels.values():
            averages['alpha_l'] += channel.alpha_l / n_channels
            averages['alpha_v'] += channel.alpha_v / n_channels
            averages['u_l'] += channel.u_l / n_channels
            averages['u_v'] += channel.u_v / n_channels
            averages['T_l'] += channel.T_l / n_channels
            averages['T_v'] += channel.T_v / n_channels
            averages['pressure'] += channel.pressure / n_channels
            averages['T_w'] += channel.T_w / n_channels
            
        return averages
        
    @staticmethod
    def calculate_mass_flow_rates(channels):
        """
        计算各个子通道的质量流量
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        dict, 包含各个子通道质量流量的字典
        """
        mass_flows = {}
        
        for channel_id, channel in channels.items():
            # 计算入口质量流量
            props = channel.fluid_props.get_properties(
                channel.pressure[0], channel.T_l[0], 'liquid')
            if props is None:
                continue
                
            m_in = props['density'] * channel.u_l[0] * channel.area
            
            # 计算出口质量流量
            props_l = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_l[-1], 'liquid')
            props_v = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_v[-1], 'vapor')
                
            if props_l is None or props_v is None:
                continue
                
            m_out = (channel.alpha_l[-1] * props_l['density'] * channel.u_l[-1] +
                    channel.alpha_v[-1] * props_v['density'] * channel.u_v[-1]) * channel.area
            
            mass_flows[channel_id] = {
                'inlet': m_in,
                'outlet': m_out,
                'difference': m_out - m_in
            }
            
        return mass_flows
        
    @staticmethod
    def calculate_energy_balance(channels):
        """
        计算能量平衡
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        dict, 包含能量平衡结果的字典
        """
        energy_balance = {}
        
        for channel_id, channel in channels.items():
            # 计算入口能量流率
            props = channel.fluid_props.get_properties(
                channel.pressure[0], channel.T_l[0], 'liquid')
            if props is None:
                continue
                
            e_in = props['density'] * channel.u_l[0] * props['enthalpy'] * channel.area
            
            # 计算出口能量流率
            props_l = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_l[-1], 'liquid')
            props_v = channel.fluid_props.get_properties(
                channel.pressure[-1], channel.T_v[-1], 'vapor')
                
            if props_l is None or props_v is None:
                continue
                
            e_out = (channel.alpha_l[-1] * props_l['density'] * channel.u_l[-1] * 
                    props_l['enthalpy'] +
                    channel.alpha_v[-1] * props_v['density'] * channel.u_v[-1] * 
                    props_v['enthalpy']) * channel.area
            
            # 计算加热功率
            q_total = cfg.MAX_HEAT_FLUX * cfg.CHANNEL_LENGTH
            
            energy_balance[channel_id] = {
                'inlet': e_in,
                'outlet': e_out,
                'heat_added': q_total,
                'imbalance': e_out - e_in - q_total
            }
            
        return energy_balance
        
    @staticmethod
    def calculate_pressure_drop(channels):
        """
        计算压降
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        dict, 包含压降结果的字典
        """
        pressure_drops = {}
        
        for channel_id, channel in channels.items():
            # 计算总压降
            total_dp = channel.pressure[0] - channel.pressure[-1]
            
            # 计算重力压降
            props = channel.fluid_props.get_properties(
                channel.pressure[0], channel.T_l[0], 'liquid')
            if props is None:
                continue
                
            gravity_dp = props['density'] * 9.81 * cfg.CHANNEL_LENGTH
            
            # 计算摩擦压降（简化处理）
            friction_dp = total_dp - gravity_dp
            
            pressure_drops[channel_id] = {
                'total': total_dp,
                'gravity': gravity_dp,
                'friction': friction_dp
            }
            
        return pressure_drops
        
    @staticmethod
    def calculate_heat_transfer_coefficients(channels):
        """
        计算传热系数
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        dict, 包含传热系数的字典
        """
        htc = {}
        
        for channel_id, channel in channels.items():
            htc[channel_id] = np.zeros(cfg.AXIAL_NODES)
            
            for i in range(cfg.AXIAL_NODES):
                # 计算壁面热流密度
                q_wall = cfg.MAX_HEAT_FLUX / (np.pi * channel.diameter)
                
                # 计算温差
                delta_T = channel.T_w[i] - channel.T_l[i]
                
                # 计算传热系数
                if abs(delta_T) > 1e-6:  # 避免除以零
                    htc[channel_id][i] = q_wall / delta_T
                else:
                    htc[channel_id][i] = 0.0
                    
        return htc
        
    @staticmethod
    def export_results(channels, filename):
        """
        导出计算结果到文件
        
        参数：
        channels: dict, 子通道字典
        filename: str, 文件名
        """
        # 计算轴向位置
        z = np.linspace(0, cfg.CHANNEL_LENGTH, cfg.AXIAL_NODES)
        
        # 打开文件
        with open(filename, 'w') as f:
            # 写入表头
            f.write('Axial_Position(m),Channel_ID,Alpha_l,Alpha_v,U_l(m/s),U_v(m/s),' +
                   'T_l(K),T_v(K),P(Pa),T_w(K)\n')
            
            # 写入数据
            for i in range(cfg.AXIAL_NODES):
                for channel_id, channel in channels.items():
                    f.write(f'{z[i]:.6f},{channel_id},' +
                           f'{channel.alpha_l[i]:.6f},{channel.alpha_v[i]:.6f},' +
                           f'{channel.u_l[i]:.6f},{channel.u_v[i]:.6f},' +
                           f'{channel.T_l[i]:.6f},{channel.T_v[i]:.6f},' +
                           f'{channel.pressure[i]:.6f},{channel.T_w[i]:.6f}\n')
                           
    @staticmethod
    def calculate_statistics(channels):
        """
        计算统计数据
        
        参数：
        channels: dict, 子通道字典
        
        返回：
        dict, 包含统计结果的字典
        """
        stats = {}
        
        # 获取所有通道的数据
        all_alpha_l = np.array([channel.alpha_l for channel in channels.values()])
        all_alpha_v = np.array([channel.alpha_v for channel in channels.values()])
        all_u_l = np.array([channel.u_l for channel in channels.values()])
        all_u_v = np.array([channel.u_v for channel in channels.values()])
        all_T_l = np.array([channel.T_l for channel in channels.values()])
        all_T_v = np.array([channel.T_v for channel in channels.values()])
        all_p = np.array([channel.pressure for channel in channels.values()])
        all_T_w = np.array([channel.T_w for channel in channels.values()])
        
        # 计算统计量
        stats['alpha_l'] = {
            'mean': np.mean(all_alpha_l),
            'std': np.std(all_alpha_l),
            'min': np.min(all_alpha_l),
            'max': np.max(all_alpha_l)
        }
        
        stats['alpha_v'] = {
            'mean': np.mean(all_alpha_v),
            'std': np.std(all_alpha_v),
            'min': np.min(all_alpha_v),
            'max': np.max(all_alpha_v)
        }
        
        stats['u_l'] = {
            'mean': np.mean(all_u_l),
            'std': np.std(all_u_l),
            'min': np.min(all_u_l),
            'max': np.max(all_u_l)
        }
        
        stats['u_v'] = {
            'mean': np.mean(all_u_v),
            'std': np.std(all_u_v),
            'min': np.min(all_u_v),
            'max': np.max(all_u_v)
        }
        
        stats['T_l'] = {
            'mean': np.mean(all_T_l),
            'std': np.std(all_T_l),
            'min': np.min(all_T_l),
            'max': np.max(all_T_l)
        }
        
        stats['T_v'] = {
            'mean': np.mean(all_T_v),
            'std': np.std(all_T_v),
            'min': np.min(all_T_v),
            'max': np.max(all_T_v)
        }
        
        stats['pressure'] = {
            'mean': np.mean(all_p),
            'std': np.std(all_p),
            'min': np.min(all_p),
            'max': np.max(all_p)
        }
        
        stats['T_w'] = {
            'mean': np.mean(all_T_w),
            'std': np.std(all_T_w),
            'min': np.min(all_T_w),
            'max': np.max(all_T_w)
        }
        
        return stats 